{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7e7d9baf-a5f6-4aee-a150-734bcaf415e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import glob\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "275c57c4-f47c-4d05-8899-756971622bad",
   "metadata": {},
   "outputs": [],
   "source": [
    "config_filename = \"sppin_config.json\"\n",
    "\n",
    "config = dict()\n",
    "\n",
    "model_config = dict()\n",
    "model_config[\"name\"] = \"DynUNet\"  # network model name from MONAI\n",
    "# set the network hyper-parameters\n",
    "model_config[\"in_channels\"] = 4  # 4 input images for the BraTS challenge\n",
    "model_config[\"out_channels\"] = 1   # whole tumor, tumor core, enhancing tumor\n",
    "model_config[\"spatial_dims\"] = 3   # 3D input images\n",
    "model_config[\"deep_supervision\"] = False  # do not check outputs of lower layers\n",
    "model_config[\"strides\"] = [[1, 1, 1], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]][:-1]  # number of downsampling convolutions\n",
    "model_config[\"filters\"] = [64, 96, 128, 192, 256, 384, 512, 768, 1024][:len(model_config[\"strides\"])]  # number of filters per layer\n",
    "model_config[\"kernel_size\"] = [[3, 3, 3]] * len(model_config[\"strides\"])  # size of the convolution kernels per layer\n",
    "model_config[\"upsample_kernel_size\"] = model_config[\"strides\"][1:]  # should be the same as the strides\n",
    "\n",
    "# put the model config in the main config\n",
    "config[\"model\"] = model_config\n",
    "\n",
    "config[\"optimizer\"] = {'name': 'Adam', \n",
    "                       'lr': 0.001}  # initial learning rate\n",
    "\n",
    "# define the loss\n",
    "config[\"loss\"] = {'name': 'DiceLoss', # from Monai\n",
    "                  'include_background': True,  # we do not have a label for the background, so this should be true (by \"include background\" monai means include channel 0)\n",
    "                  'sigmoid': True,  # transform the model logits to activations\n",
    "                  'batch': False}  \n",
    "\n",
    "# set the cross validation parameters\n",
    "config[\"cross_validation\"] = {'folds': 5,  # number of cross validation folds\n",
    "                              'seed': 25},  # seed to make the generation of cross validation folds consistent across different trials\n",
    "# set the scheduler parameters\n",
    "config[\"scheduler\"] = {'name': 'ReduceLROnPlateau', \n",
    "                       'patience': 20,  # wait 10 epochs with no improvement before reducing the learning rate\n",
    "                       'factor': 0.5,   # multiply the learning rate by 0.5\n",
    "                       'min_lr': 1e-08}  # stop reducing the learning rate once it gets to 1e-8\n",
    "\n",
    "# set the dataset parameters\n",
    "config[\"dataset\"] = {'name': 'SegmentationDatasetPersistent',  # 'Persistent' means that it will save the preprocessed outputs generated during the first epoch\n",
    "# However, using 'Persistent', does also increase the time of the first epoch compared to the other epochs, which should run faster\n",
    "  'desired_shape': [192, 192, 192],  # resize the images to this shape, increase this to get higher resolution images (increases computation time and memory usage)\n",
    "  'labels': [1],  # 1: tumor\n",
    "  'orientation': 'RAS',  # Force all the images to be the same orientation (Right-Anterior-Suppine)\n",
    "  'normalization': 'NormalizeIntensityD',  # z score normalize the input images to zero mean unit standard deviation\n",
    "  'normalization_kwargs': {'channel_wise': True, \"nonzero\": False},  # perform the normalization channel wise and include the background\n",
    "  'resample': True,  # resample the images when resizing them, otherwise the resize could crop out regions of interest\n",
    "  'crop_foreground': True,  # crop the foreground of the images\n",
    "  'foreground_percentile': 0.9,  # aggressive foreground cropping to make sure the empty space is taken out of the images\n",
    "  'training':  # the following arguments will only be applied to the training data.\n",
    "    {\n",
    "    'spatial_augmentations': [{'name': 'RandFlipD', 'spatial_axis': 0, 'prob': 0.5},\n",
    "                              {'name': 'RandFlipD', 'spatial_axis': 1, 'prob': 0.5},\n",
    "                              {'name': 'RandRotateD', 'prob': 0.5, 'range_x': 0.2, 'range_y': 0.2, 'range_z': 0.2}],\n",
    "    'intensity_augmentations': [{'name': 'RandScaleIntensityD', 'factors': 0.1, 'prob': 1.0},\n",
    "                                {'name': 'RandShiftIntensityD', 'offsets': 0.1, 'prob': 1.0}],\n",
    "    }\n",
    "                    }\n",
    "config[\"training\"] = {'batch_size': 2,  # number of image/label pairs to read at a time during training\n",
    "  'validation_batch_size': 2,  # number of image/label pairs to read at atime during validation\n",
    "  'amp': False,  # don't set this to true unless the model you are using is setup to use automatic mixed precision (AMP)\n",
    "  'early_stopping_patience': None,  # stop the model early if the validaiton loss stops improving\n",
    "  'n_epochs': 1000,  # number of training epochs, reduce this if you don't want training to run as long\n",
    "  'save_every_n_epochs': None,  # save the model every n epochs (otherwise only the latest model will be saved)\n",
    "  'save_last_n_models': None,  # save the last n models \n",
    "  'save_best': True}  # save the model that has the best validation loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "652520c3-0cb7-42a6-9392-98f2776850cc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['dwi_b0', 'dwi_b100', 'nb', 't1_gd', '']\n"
     ]
    }
   ],
   "source": [
    "# get the training filenames\n",
    "config[\"training_filenames\"] = list()\n",
    "ground_truth_filenames = sorted(glob.glob(\"./aligned/*/*/*NB*.nii*\"))\n",
    "for label_filename in ground_truth_filenames:\n",
    "    subject, visit = label_filename.split(\"/\")[-3:-1]\n",
    "    filenames = [os.path.abspath(fn) for fn in sorted(glob.glob(os.path.join(os.path.dirname(label_filename), \"*\")))]\n",
    "    n_features = len(filenames) \n",
    "    feature_modalities = [\"_\".join(fn.split(\"_\")[3:-1]).strip(\"8 \").lower() for fn in filenames]\n",
    "    t1_filename = filenames[feature_modalities.index(\"T1_gd\".lower())]\n",
    "    assert os.path.exists(t1_filename)\n",
    "    \n",
    "    if len(feature_modalities) < 5:\n",
    "        continue\n",
    "    \n",
    "    if \"T2\".lower() not in feature_modalities:\n",
    "        for filename in filenames:\n",
    "            if \"T2\" in filename:\n",
    "                t2_fn = filename\n",
    "        # print(t2_fn)\n",
    "        print(feature_modalities)\n",
    "    else:\n",
    "        t2_fn = filenames[feature_modalities.index(\"T2\".lower())]\n",
    "    assert os.path.exists(t2_fn)\n",
    "    \n",
    "    if \"DWI_b0\".lower() not in feature_modalities:\n",
    "        for filename in filenames:\n",
    "            if \"DWI\" in filename and \"b0\" in filename:\n",
    "                dwi_b0 = filename\n",
    "        # print(dwi_b0)\n",
    "    else:\n",
    "        dwi_b0 = filenames[feature_modalities.index(\"DWI_b0\".lower())]\n",
    "    assert os.path.exists(dwi_b0)\n",
    "\n",
    "    if \"DWI_b100\".lower() not in feature_modalities:\n",
    "        for filename in filenames:\n",
    "            if \"DWI\" in filename and \"b100\" in filename:\n",
    "                dwi_b100 = filename\n",
    "        # print(dwi_b100)\n",
    "    else:\n",
    "        dwi_b100 = filenames[feature_modalities.index(\"DWI_b100\".lower())]\n",
    "    assert os.path.exists(dwi_b100)\n",
    "\n",
    "    # print(t2_fn)\n",
    "    config[\"training_filenames\"].append({\"image\": [t1_filename, t2_fn, dwi_b0, dwi_b100], \"label\": label_filename})\n",
    "with open(config_filename, \"w\") as op:\n",
    "    json.dump(config, op, indent=4)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (brats)",
   "language": "python",
   "name": "brats"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
